{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the TensorBoard notebook extension.\n",
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorFlow version:  2.1.0\n"
     ]
    }
   ],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "from datetime import datetime\n",
    "from packaging import version\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import Model, layers, activations, Sequential\n",
    "\n",
    "print(\"TensorFlow version: \", tf.__version__)\n",
    "assert version.parse(tf.__version__).release[0] >= 2, \"This notebook requires TensorFlow 2.0 or above.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!rd /S /Q logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "# Reference\n",
    "\n",
    "[Densely Connected Convolutional Networks]()\n",
    "\n",
    "This model is based on the following implementations:\n",
    "- https://github.com/calmisential/DenseNet_TensorFlow2/blob/master/densenet.py\n",
    "\n",
    "\"\"\"\n",
    "class ConvBlock(layers.Layer):\n",
    "    def __init__(self, growth_rate=32):\n",
    "        super(ConvBlock, self).__init__()\n",
    "        self.block_bn1 = layers.BatchNormalization()\n",
    "        self.block_relu1 = layers.ReLU()\n",
    "        self.block_conv1 = layers.Conv2D(4 * growth_rate, 1, 1, \"same\")\n",
    "\n",
    "        self.block_bn2 = layers.BatchNormalization()\n",
    "        self.block_relu2 = layers.ReLU()\n",
    "        self.block_conv2 = layers.Conv2D(growth_rate, 3, 1, \"same\")\n",
    "\n",
    "        self.block_concat = layers.Concatenate(axis=3)\n",
    "\n",
    "    def call(self, inputs, training=None, **kwargs):\n",
    "        x = self.block_bn1(inputs, training=training)\n",
    "        x = self.block_relu1(x)\n",
    "        x = self.block_conv1(x)\n",
    "\n",
    "        x = self.block_bn2(x, training=training)\n",
    "        x = self.block_relu2(x)\n",
    "        x = self.block_conv2(x)\n",
    "\n",
    "        x = self.block_concat([inputs, x])\n",
    "        return x\n",
    "\n",
    "\n",
    "class DenseBlock(layers.Layer):\n",
    "    def __init__(self, num_layers, growth_rate):\n",
    "        super(DenseBlock, self).__init__()\n",
    "        self.num_layers = num_layers\n",
    "        self.growth_rate = growth_rate\n",
    "        self.dense_layer = self._make_layer()\n",
    "\n",
    "    def _make_layer(self):\n",
    "        features_list = []\n",
    "        for i in range(self.num_layers):\n",
    "            features_list.append(ConvBlock(self.growth_rate))\n",
    "        return Sequential(features_list)\n",
    "\n",
    "    def call(self, inputs, training=None, **kwargs):\n",
    "        x = self.dense_layer(inputs)\n",
    "        return x\n",
    "\n",
    "\n",
    "class TransitionLayer(layers.Layer):\n",
    "    def __init__(self, out_channels):\n",
    "        super(TransitionLayer, self).__init__()\n",
    "        self.trans_bn1 = layers.BatchNormalization()\n",
    "        self.trans_relu1 = layers.ReLU()\n",
    "        self.trans_conv1 = layers.Conv2D(out_channels, 1, 1, \"same\")\n",
    "        self.trans_pool1 = layers.AveragePooling2D(2, 2, \"same\")\n",
    "\n",
    "    def call(self, inputs, training=None, **kwargs):\n",
    "        x = self.trans_bn1(inputs)\n",
    "        x = self.trans_relu1(x)\n",
    "        x = self.trans_conv1(x)\n",
    "        x = self.trans_pool1(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class DenseNet(Model):\n",
    "    def __init__(self, num_init_features, growth_rate, block_layers, compression_rate, drop_rate, num_class=200):\n",
    "        super(DenseNet, self).__init__()\n",
    "        self.conv1 = layers.Conv2D(num_init_features, 7, 2, \"same\")\n",
    "        self.bn1 = layers.BatchNormalization()\n",
    "        self.relu1 = layers.ReLU()\n",
    "        self.pool1 = layers.MaxPool2D(3, 2, \"same\")\n",
    "\n",
    "        self.num_channels = num_init_features\n",
    "        self.dense_block_1 = DenseBlock(num_layers=block_layers[0], growth_rate=growth_rate)\n",
    "        self.num_channels += growth_rate * block_layers[0]\n",
    "        self.num_channels = compression_rate * self.num_channels\n",
    "        self.transition_1 = TransitionLayer(out_channels=int(self.num_channels))\n",
    "\n",
    "        self.dense_block_2 = DenseBlock(num_layers=block_layers[1], growth_rate=growth_rate)\n",
    "        self.num_channels += growth_rate * block_layers[1]\n",
    "        self.num_channels = compression_rate * self.num_channels\n",
    "        self.transition_2 = TransitionLayer(out_channels=int(self.num_channels))\n",
    "\n",
    "        self.dense_block_3 = DenseBlock(num_layers=block_layers[2], growth_rate=growth_rate)\n",
    "        self.num_channels += growth_rate * block_layers[2]\n",
    "        self.num_channels = compression_rate * self.num_channels\n",
    "        self.transition_3 = TransitionLayer(out_channels=int(self.num_channels))\n",
    "\n",
    "        self.dense_block_4 = DenseBlock(num_layers=block_layers[3], growth_rate=growth_rate)\n",
    "\n",
    "        self.bn2 = layers.BatchNormalization()\n",
    "        self.relu2 = layers.ReLU()\n",
    "        self.avg_pool = layers.GlobalAveragePooling2D()\n",
    "        self.fc = layers.Dense(units=num_class, activation=activations.softmax)\n",
    "\n",
    "    def call(self, inputs, training=None, **kwargs):\n",
    "        x = self.conv1(inputs)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu1(x)\n",
    "        x = self.pool1(x)\n",
    "\n",
    "        x = self.dense_block_1(x, training=training)\n",
    "        x = self.transition_1(x, training=training)\n",
    "        x = self.dense_block_2(x, training=training)\n",
    "        x = self.transition_2(x, training=training)\n",
    "        x = self.dense_block_3(x, training=training)\n",
    "        x = self.transition_3(x, training=training)\n",
    "        x = self.dense_block_4(x, training=training)\n",
    "\n",
    "        x = self.bn2(x)\n",
    "        x = self.relu2(x)\n",
    "        x = self.avg_pool(x)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "def densenet121():\n",
    "    return DenseNet(64, 32, [6, 12, 24, 16], 0.5, 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = densenet121()\n",
    "model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((60000, 28, 28, 1), (60000,))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load and prepare the MNIST dataset.\n",
    "\n",
    "mnist = tf.keras.datasets.mnist\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train, x_test = x_train / 255.0, x_test / 255.0\n",
    "\n",
    "# # Add a channels dimension\n",
    "x_train = x_train[..., tf.newaxis]\n",
    "x_test  = x_test[..., tf.newaxis]\n",
    "x_train.shape, y_train.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 60000 samples\n",
      "60000/60000 [==============================] - 293s 5ms/sample - loss: 0.1448 - accuracy: 0.9586\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x159b540db48>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Define the Keras TensorBoard callback.\n",
    "logdir=\"logs\\\\fit\\\\\" + datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)\n",
    "\n",
    "# Train the model.\n",
    "model.fit(x_train, y_train, batch_size=64, epochs=1, callbacks=[tensorboard_callback])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
